{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ablating neurons in a SPA model\n",
    "\n",
    "The model used here is a SPA model that goes through\n",
    "a set sequence of semantic pointers,\n",
    "as shown in\n",
    "[this core Nengo example](https://www.nengo.ai/nengo/examples/spa_sequence.html).\n",
    "\n",
    "This is primarily a demonstration of how to apply the `ablate_ensemble` function\n",
    "to all the ensembles in a network or SPA module.\n",
    "If you haven't gone through the `ablate_ensemble` example\n",
    "in this directory, go through that example first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import nengo\n",
    "%load_ext nengo.ipynb\n",
    "from nengo import spa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dimensions = 16\n",
    "\n",
    "def start(t):\n",
    "    if t < 0.05:\n",
    "        return 'A'\n",
    "    else:\n",
    "        return '0'\n",
    "\n",
    "with spa.SPA() as model:\n",
    "    model.cortex = spa.Buffer(dimensions=dimensions)\n",
    "    actions = spa.Actions(\n",
    "        'dot(cortex, A) --> cortex = B',\n",
    "        'dot(cortex, B) --> cortex = C',\n",
    "        'dot(cortex, C) --> cortex = D',\n",
    "        'dot(cortex, D) --> cortex = E',\n",
    "        'dot(cortex, E) --> cortex = A'\n",
    "    )\n",
    "    model.bg = spa.BasalGanglia(actions=actions)\n",
    "    model.thal = spa.Thalamus(model.bg)\n",
    "    model.input = spa.Input(cortex=start)\n",
    "\n",
    "    cortex = nengo.Probe(model.cortex.state.output, synapse=0.01)\n",
    "    actions = nengo.Probe(model.thal.actions.output, synapse=0.01)\n",
    "    utility = nengo.Probe(model.bg.input, synapse=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    sim.run(0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def plot():\n",
    "    fig = plt.figure(figsize=(12,8))\n",
    "    p1 = fig.add_subplot(3,1,1)\n",
    "\n",
    "    p1.plot(sim.trange(), model.similarity(sim.data, cortex))\n",
    "    p1.legend(model.get_output_vocab('cortex').keys, fontsize='x-small')\n",
    "    p1.set_ylabel('State')\n",
    "\n",
    "    p2 = fig.add_subplot(3,1,2)\n",
    "    p2.plot(sim.trange(), sim.data[actions])\n",
    "    p2_legend_txt = [a.effect for a in model.bg.actions.actions]\n",
    "    p2.legend(p2_legend_txt, fontsize='x-small')\n",
    "    p2.set_ylabel('Action')\n",
    "\n",
    "    p3 = fig.add_subplot(3,1,3)\n",
    "    p3.plot(sim.trange(), sim.data[utility])\n",
    "    p3_legend_txt = [a.condition for a in model.bg.actions.actions]\n",
    "    p3.legend(p3_legend_txt, fontsize='x-small')\n",
    "    p3.set_ylabel('Utility')\n",
    "\n",
    "    fig.subplots_adjust(hspace=0.2)\n",
    "plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the same `ablate_ensemble` function as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def ablate_ensemble(ens, proportion, sim, bias=True):\n",
    "    \"\"\"Ablate a proportion of the neurons in an ensemble.\n",
    "\n",
    "    The ablation is done by setting the encoder and gain associated\n",
    "    with a neuron to zero. Since no input current being injected,\n",
    "    the neuron will generally be silent. However, if there is direct\n",
    "    current injected with a neuron-to-neuron connection, then the\n",
    "    cell may still fire. To counter that in most cases, we set the\n",
    "    bias associated with the neuron to a large negative value.\n",
    "    \"\"\"\n",
    "\n",
    "    n_neurons = min(int(ens.n_neurons * proportion), ens.n_neurons)\n",
    "    idx = np.random.choice(np.arange(ens.n_neurons), replace=False, size=n_neurons)\n",
    "\n",
    "    encoder_sig = sim.signals[sim.model.sig[ens]['encoders']]\n",
    "    encoder_sig.setflags(write=True)\n",
    "    encoder_sig[idx] = 0.0\n",
    "    encoder_sig.setflags(write=False)\n",
    "\n",
    "    if bias:\n",
    "        bias_sig = sim.signals[sim.model.sig[ens.neurons]['bias']]\n",
    "        bias_sig.setflags(write=True)\n",
    "        bias_sig[idx] = -1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, however, instead of applying it to a single ensemble,\n",
    "we will apply it to all of the ensembles in a SPA module.\n",
    "\n",
    "We can therefore see the effects of ablation in different\n",
    "parts of the model. Let's see what happens when ablating\n",
    "25% of the neurons in the cortex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    for ens in model.cortex.all_ensembles:\n",
    "        ablate_ensemble(ens, 0.25, sim)\n",
    "    sim.run(0.5)\n",
    "plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In some cases, the sequence stops after some time, which is interesting!\n",
    "\n",
    "It saves us a tiny bit of typing to make a helper function\n",
    "to ablate whole networks (or SPA modules).\n",
    "So let's do that, then look at the effects of ablating\n",
    "the other parts of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def ablate_network(net, proportion, sim, bias=True):\n",
    "    \"\"\"Ablate a proportion of the neurons in all ensembles in a network.\"\"\"\n",
    "    for ens in net.all_ensembles:\n",
    "        ablate_ensemble(ens, proportion, sim, bias=bias)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    ablate_network(model.bg, 0.25, sim)\n",
    "    sim.run(0.5)\n",
    "plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    ablate_network(model.thal, 0.25, sim)\n",
    "    sim.run(0.5)\n",
    "plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's look at what happens when we ablate\n",
    "25% of the neurons in the entire model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    ablate_network(model, 0.25, sim)\n",
    "    sim.run(0.5)\n",
    "plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not silencing the neurons seems to hurt performance,\n",
    "possibly due to the added noise of the random spikes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    ablate_network(model, 0.25, sim, bias=False)\n",
    "    sim.run(0.5)\n",
    "plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To verify that this is indeed working as expected,\n",
    "we can ablate all of the neurons and confirm that\n",
    "there is no activity in the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:\n",
    "    ablate_network(model, 1, sim)\n",
    "    sim.run(0.5)\n",
    "plot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  },
  "widgets": {
   "state": {
    "00013f0ca1bc40658181b8438fa8cf5e": {
     "views": []
    },
    "01794e20a9cc485b974d27a1d207fa2b": {
     "views": []
    },
    "06e22a5624aa400b841f59826be7ac93": {
     "views": []
    },
    "08f7eb9342ea4ed0b95bbd962cf6985b": {
     "views": []
    },
    "098bf839bf4442f19d0f05423e0da92c": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    },
    "0c0b025970f344a89f0183770ab8a89f": {
     "views": []
    },
    "0fd020ee43564878aa0e102f7933ba9a": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "1018ba8bab444291891da85b4bbb3f85": {
     "views": []
    },
    "136194de0f5e4f92b822821043ac6c24": {
     "views": []
    },
    "1922c4ca3c0e4cbfbac7401a109ccd46": {
     "views": []
    },
    "19fe9288d57946e58a66e542b11de33c": {
     "views": []
    },
    "1bd5e8a952bf443f87c7d7352597416f": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    },
    "1dcc4b5d12904338845bbda398a0805c": {
     "views": []
    },
    "24b4992658224412ab268fc406d0741d": {
     "views": [
      {
       "cell_index": 8
      }
     ]
    },
    "2a139eeeee194143ab1741a36e3692c5": {
     "views": [
      {
       "cell_index": 15
      }
     ]
    },
    "2b0dc0b80b0d488fbaa93553b4fe539e": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    },
    "2b11df95ba234ee8b63e7a41fed88919": {
     "views": []
    },
    "2b12902053a84f1688c50af9d9d98690": {
     "views": []
    },
    "3352ee992a4440528a481755d8fdbad3": {
     "views": []
    },
    "35a6d22f200848f98aafd7c36b050be0": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "36fff258a29e48c4b4db392ade7f8784": {
     "views": []
    },
    "394e7250d24249cd9d360f518f5ca3d5": {
     "views": [
      {
       "cell_index": 3
      }
     ]
    },
    "3b17ce9db3994654ba61791bfa312ae1": {
     "views": [
      {
       "cell_index": 12
      }
     ]
    },
    "3d92442d181c4b519357f02db5fb8943": {
     "views": [
      {
       "cell_index": 11
      }
     ]
    },
    "3e761d874105461c9f8819de431605e3": {
     "views": []
    },
    "407d23c74d514129bc64f2794542cc78": {
     "views": [
      {
       "cell_index": 8
      }
     ]
    },
    "40a5d5a4cd2044588f2204dd1a3f852d": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    },
    "46408b8ef39940f6af415dff70ab7dbc": {
     "views": [
      {
       "cell_index": 11
      }
     ]
    },
    "49afbebd7a234295b3f85988169fc39f": {
     "views": []
    },
    "51203848844a4319a625815ed4a1f805": {
     "views": []
    },
    "5315148967bb4ea49aa08fb54377ddcc": {
     "views": []
    },
    "558b9a55ef3a4642ac28191d941ae389": {
     "views": []
    },
    "63b0d0d6d5dc4575a728d96bc80cdfdf": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "67c190e7fe6d45559cb1794b96a9ed82": {
     "views": [
      {
       "cell_index": 14
      }
     ]
    },
    "7919655309384f17a89e44a852b83d72": {
     "views": [
      {
       "cell_index": 15
      }
     ]
    },
    "79c0df1a382446ef9d06b3c78e2c7531": {
     "views": []
    },
    "85eed3febe694ea3ad0806e89bc91920": {
     "views": []
    },
    "8976b6c4edc240e2aea67dbc2bd2b3f1": {
     "views": [
      {
       "cell_index": 3
      }
     ]
    },
    "8d06140aa75f494087ca896a48f40b35": {
     "views": []
    },
    "9503fa7ebf74426284c7a4613518094d": {
     "views": []
    },
    "95399f9749704d5f99ba04f63fa5b1f2": {
     "views": []
    },
    "97d27c09966e43499f9ddfcfa68920d9": {
     "views": [
      {
       "cell_index": 12
      }
     ]
    },
    "a7908f92e1244f1d850fac1ff3813196": {
     "views": []
    },
    "abe1f008754f4372ba3c93a9daf0d1dd": {
     "views": []
    },
    "ac4cad8ebc30465881953622d59e75bd": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    },
    "ad86ac7242834c1fada189bcdcd7ac98": {
     "views": [
      {
       "cell_index": 8
      }
     ]
    },
    "af0f31d0ace44a96a3095d171087cef1": {
     "views": []
    },
    "b114d87d2c254efd886efffe533d5421": {
     "views": []
    },
    "b12b734fd02e43a6a96077edf91a6918": {
     "views": [
      {
       "cell_index": 11
      }
     ]
    },
    "b29df4a79ad8428a8995203a6edd4d17": {
     "views": []
    },
    "b795904f81144cf59fdc0456bafd213e": {
     "views": []
    },
    "b84ed054ed1a47748e35672867c7a651": {
     "views": [
      {
       "cell_index": 16
      }
     ]
    },
    "bd905af700a649c792e7555464efc38a": {
     "views": []
    },
    "bea6009e0ecb452497defa49ab31c9fa": {
     "views": []
    },
    "bfe9969bd64b413296159be1fda07b46": {
     "views": []
    },
    "c09cb0a3efcd4d65abc1eccb5e3ea70e": {
     "views": []
    },
    "c39ad9e7d1a24bad8fe41a15608b03aa": {
     "views": []
    },
    "c655fc3b44c24a619a7c347fdddfafc8": {
     "views": []
    },
    "c76bb64489fc4dfda34dd27bda92c33d": {
     "views": []
    },
    "cf419753cb514344af25d8d8f487925f": {
     "views": []
    },
    "d3d17cb909074fd2bbc4860f70f31c3e": {
     "views": [
      {
       "cell_index": 11
      }
     ]
    },
    "d4cf26823f6846fbb6a092a83ae6832d": {
     "views": []
    },
    "e2b4b552dd124bf3964f14f3680ee4ad": {
     "views": []
    },
    "f1730fae9ab6437fad56cccaac47e9e1": {
     "views": []
    },
    "f6c59c84942a4c348c0f9f92ff27e6c4": {
     "views": []
    },
    "fd48f617336e4007ad0293bcc23839db": {
     "views": []
    }
   },
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
